# 这个币开始锁仓还没多久，一些参数可能还存在错误，无所谓，随时可以对比更正
# 账号模型
class RModule:
    """
    args:直接下线实例
    """
    def __init__(self, *args):
        self.A = 0      # 可用
        self.B = 0      # 令牌
        self.D = 0      # 自主
        self.left = 0   # 剩余出局令牌数量，3倍锁入
        self.left_flag = 0  # 标记是否要重置self.left
        self.putin = 0  # 本次投入的合约数量
        self.earnings = 0   # 收益
        self.day = 0    # 天数调整
        self.ls = list(args)  # 实例列表，考虑可以后续新增
        self.income1 = 0   # 合约收益
        self.income2 = 0   # 分享增量
        self.income3 = 0   # 管理增量
        self.hl = 0        # 红利投入
        self.hl6 = 0       # 6倍红利

    # 锁入合约
    def lock_coin(self, coin, day=0):
        self.putin = coin
        self.day = day
    
    # 锁入红利
    def lock_hl(self, coin, flag=1):
        if flag:
            self.A -= coin*3/5
            self.D -= coin*2/5
        else:
            self.A -= coin
        self.hl += coin
        self.hl6 += 6*coin

    # 每天收益情况
    def one_day(self, day):
        a, b = self.A, self.B
        # 合约收益
        self.A += round(self.putin*(level_dic[self]*0.1+0.4+0.1*(day-self.day))/100*0.7,4)
        self.B += round(self.putin*(level_dic[self]*0.1+0.4+0.1*(day-self.day))/100*0.3,4)
        self.left -= round(self.putin*(level_dic[self]*0.1+0.4+0.1*(day-self.day))/100*0.3,4)
        if self.left<=0:
            self.left += self.putin*3   # 3倍出局
        self.earnings = self.A + self.B - a - b
        self.income1 += round(self.putin*(level_dic[self]*0.1+0.4+0.1*(day-self.day))/100,4)
        # 红利收益
        if self.hl6 > 0:
            self.B += self.hl * 0.002 * 0.02
            self.D += self.hl * 0.002 * 0.73
            self.hl6 -= self.hl * 0.002
            self.left -= self.hl * 0.002 * 0.02
        # 分享收益+管理收益
        rate2 = min(1, len(self.ls)/10)
        # rate2 = sum([i.coin>0 for i in self.ls])/10   # 分享收益率似乎并不看下线是否开合约，数据少暂不确定
        for i in self.ls:
            i.one_day(day)   # 下线一天后的账户情况
            # 直接下线收益(分享收益)
            self.B += round(i.earnings*rate2*0.02,4)
            self.D += round(i.earnings*rate2*0.73,4)
            self.left -= round(i.earnings*rate2*0.02,4)
            self.income2 += round(i.earnings*rate2*0.75,4)
            # 间接下线收益(管理收益)
            try:
                self.B += round(sum([j.earnings for j in i.ls])*(0.24-level_dic[self]*0.03)*0.02,4)  # 遍历第2层下线
                self.D += round(sum([j.earnings for j in i.ls])*(0.24-level_dic[self]*0.03)*0.73,4)
                self.left -= round(sum([j.earnings for j in i.ls])*(0.24-level_dic[self]*0.03)*0.02,4)
                self.income3 += round(sum([j.earnings for j in i.ls])*(0.24-level_dic[self]*0.03)*0.75,4)
            except AttributeError:
                pass
    
# 复投方法
def loop(name):
    if name.A>name.D*1.5:   # 全部由可用部分复投
        name.putin = name.A//50*50
        name.A -= name.A//50*50
    else:       # 可用60%，自主40%复投
        name.putin = name.A*5/3//50*50
        name.D -= name.A*5/3//50*50*0.4
        name.A -= name.A*5/3//50*50*0.6

level_dic = {}      # 全局变量，存储各个号的级别
# 计算账号级别，下线有3个合格号才升1级
def level(dic, *args):
    for i in args:
        dic[i] = 1
        value = []
        for j in i.ls:
            try:
                value.append(dic[j])
            except KeyError:
                level(j)
        if len(value) >= 3:
            value = sorted(value, reverse=True)
            dic[i] = min(value[2]+1, 5)
    return dic

# 下面多叉树也是数据结构才看没几天刚好这模型里有用到就动手写了，不成熟之处见谅
# 先构造好多叉树
def create_tree(child:list, my_dic:dict, p_f=0, p_p=0):
    '''
    floor:层级数
    child:各层下级数列表
    p_f:当前层数
    p_p:当前同级排位
    '''
    floor = len(child)
    if p_f>=floor:
        return {}
    if not p_f:
        my_dic['r%s%s'%(p_f,p_p)]={}
    for i in range(child[p_f]):
        my_dic['r%s%s'%(p_f,p_p)]['r%s%s'%(p_f+1,i)]={}
        create_tree(child, my_dic['r%s%s'%(p_f,p_p)], p_f+1, i)
    return my_dic

# 给多叉树每个节点实例化构建实例多叉树
def create_module(tree, up_module=None):
    dic = {}
    for k in tree:
        r = RModule()
        if up_module:
            up_module.ls.append(r)
        dic[r] = create_module(tree[k], r)
    return dic

# 从实例树中取出实例进行操作
def take_module(tree, coin_ls):
    for i in tree:
        # 对每个取出的实例操作
        global module_ls, c_ls
        module_ls.append(i)
        try:
            c_ls.append(coin_ls.pop(0))
        except IndexError:
            pass
        # 迭代下去
        if tree[i]:
            take_module(tree[i], coin_ls)

# 按输入的树计算回报
def calculate_reward(arg, coin_ls=[]):
    '''
    coin_ls:按深度优先的次序分配初始资金，默认1000
    '''
    if not arg:
        print('要么输入自定义的树，要么规定各个层级的数量')
        return
    if isinstance(arg, list):
        my_tree = create_tree(arg,{})
    elif isinstance(arg, dict):
        my_tree = arg
    else:
        print('参数类型需为列表或字典')
        return
    module_tree = create_module(my_tree)
    global module_ls, level_dic, c_ls
    module_ls = []
    c_ls = []
    take_module(module_tree, coin_ls)
    level_dic = level({}, *module_ls)
    print('节点总数:',len(level_dic))
    tx_rate = []
    ky = []
    zz = []
    days = 0
    cost = 0    # 成本
    for ind, md in enumerate(module_ls):
        try:
            md.lock_coin(c_ls[ind])
            cost += c_ls[ind]
        except IndexError:
            md.lock_coin(1000)
            cost += 1000
    while 1:
        module_ls[0].one_day(days%10)
        days += 1
        # 10天本金回到可用
        if days % 10 == 0:
            for md in module_ls:
                md.A += md.putin
        # 10天复投一次，与本金回账区分开清晰点
        if days % 10 == 0:
            for md in module_ls:
                loop(md)
        # print(days)
        tx_rate.append(sum([md.B for md in module_ls]))
        ky.append(module_ls[0].A)
        zz.append(module_ls[0].D)
        # 以回本作为终止条件
        if tx_rate[-1] >= cost:
            print('回本天数',days)
            break
    # return tx_rate
    plt.plot([i/sum(c_ls) for i in tx_rate])
    plt.show()

if __name__ == '__main__':
    my_tree = {
        'a14':{
                'a11':{},
                'a12':{},
                'a13':{}
            }
    }
    calculate_reward(my_tree)
    # 这个无论是回本还是收益率、收益增长速度都高于模型1，但教训是，收益越高，潜藏风险越高